import time
from algorithm.antAlgorithm.ant import ant
from algorithm.antAlgorithm.ant import dijkstra


class Aco:
    def __init__(self, is_key, node_edge_dict, distance_matrix, start_city, end_city, num_ants, num_iterations, alpha, beta, rho, Q):
        self.key_node = []
        self.node_map = dict()
        for index, is_ in enumerate(is_key):
            if is_:
                self.node_map[len(self.key_node)] = index
                self.key_node.append(index)
        if start_city not in self.key_node:
            self.node_map[len(self.key_node)] = start_city
            self.key_node.append(start_city)
        if end_city not in self.key_node:
            self.node_map[len(self.key_node)] = end_city
            self.key_node.append(end_city)

        self.dij_map: dict
        self.dij_map = dict(dijkstra.dijkstra(self.key_node, distance_matrix))

        n = len(self.key_node)
        self.adj_matrix = [[0] * n for _ in range(n)]
        for i, u in enumerate(self.key_node):
            for j, v in enumerate(self.key_node):
                if u != v:
                    self.adj_matrix[i][j] = self.dij_map[(u, v)][0]
                else:
                    self.adj_matrix[i][j] = 0

        self.node_edge_dict = node_edge_dict
        self.start_city = start_city
        self.end_city = end_city
        self.num_ants = num_ants
        self.num_iterations = num_iterations
        self.alpha = alpha  # 信息启发因子
        self.beta = beta  # 期望启发因子
        self.rho = rho  # 信息挥发因子
        self.Q = Q  # 信息素强度

    #  主函数，蚁群算法入口，生成算法原理数据
    def work(self):
        start_time = time.time()
        start = self.key_node.index(self.start_city)
        end = self.key_node.index(self.end_city)
        best_tour, best_tour_length, displays = (
            ant.aco_tsp(self.adj_matrix, start, end, self.num_ants, self.num_iterations,
                        self.alpha, self.beta, self.rho, self.Q, 1))  # 蚁群算法核心


        elapsed_time = time.time() - start_time
        print("*****************蚁群算法********************")
        print(f"总耗时: {elapsed_time:.2f} s")
        print(f"最短路径长度: {best_tour_length}")
        print(" ".join(map(str, best_tour)))
        print("********************************************")


        origin_path = [3, 5, 0, 5, 1]
        return origin_path, displays

    def get_origin_path(self, path):
        origin_path = []
        for i in range(0, len(path) - 1):
            origin_path.extend(
                self.dij_map[(self.node_map[path[i]], self.node_map[path[i + 1]])][1][:-1]
            )
        origin_path.append(self.node_map[path[-1]])
        return origin_path

    def get_origin_path_list(self, path):
        origin_path = []
        for i in range(0, len(path) - 1):
            origin_path.append(
                self.dij_map[(self.node_map[path[i]], self.node_map[path[i + 1]])][1]
            )
        return origin_path
